package com.lucifer.cloud.boot.blog.nl2sql.controller;
import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson2.JSONObject;
import com.alibaba.nacos.shaded.com.google.common.collect.Maps;
import com.alibaba.nacos.shaded.org.checkerframework.checker.nullness.qual.Nullable;
import com.lucifer.cloud.boot.blog.nl2sql.builder.TrainBuilder;
import com.lucifer.cloud.boot.blog.nl2sql.domain.req.ChatReq;
import com.lucifer.cloud.boot.blog.nl2sql.engine.RagEngine;
import com.lucifer.cloud.boot.blog.nl2sql.engine.SpringSqlEngine;
import com.lucifer.cloud.boot.blog.nl2sql.enumd.TrainPolicyType;
import com.lucifer.cloud.boot.blog.nl2sql.vector.BaseVectorStore;
import jakarta.annotation.Resource;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.zhipuai.ZhiPuAiChatModel;
import org.springframework.core.io.ResourceLoader;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;
import reactor.core.publisher.Flux;
import reactor.core.scheduler.Schedulers;

import java.io.IOException;
import java.util.Map;


/**
 * @author lucifer
 * @date 2025/3/14 14:01
 */
@RestController
@RequestMapping("nl2sql")
public class Nl2SqlController {

    @Resource
    private BaseVectorStore vectorStore;

    @Resource
    private RagEngine ragEngine;


    @Resource
    private ZhiPuAiChatModel chatModel;


    @Resource
    private ResourceLoader resourceLoader;


    @Resource
    private SpringSqlEngine sqlEngine;


    @RequestMapping("ddl")
    public void ddl(){
        String ddl = """
                CREATE TABLE `dtp_hospital` (
                `id` BIGINT NOT NULL COMMENT '主键',
                `province` VARCHAR ( 20 ) CHARACTER SET utf8 COLLATE utf8_unicode_ci NULL DEFAULT NULL COMMENT '省份',
                `city` VARCHAR ( 20 ) CHARACTER SET utf8 COLLATE utf8_unicode_ci NULL DEFAULT NULL COMMENT '城市',
                `reporting_team` VARCHAR ( 50 ) CHARACTER SET utf8 COLLATE utf8_unicode_ci NULL DEFAULT NULL COMMENT '提报团队',
                `district` VARCHAR ( 20 ) CHARACTER SET utf8 COLLATE utf8_unicode_ci NULL DEFAULT NULL COMMENT '区',
                `hospital_name` VARCHAR ( 100 ) CHARACTER SET utf8 COLLATE utf8_unicode_ci NULL DEFAULT NULL COMMENT '申请的DTP药房 主要对应的医院名称',
                `hospital_code` VARCHAR ( 50 ) CHARACTER SET utf8 COLLATE utf8_unicode_ci NULL DEFAULT NULL COMMENT '申请DTP主要对应的医院code',
                `hospital_address` VARCHAR ( 255 ) CHARACTER SET utf8 COLLATE utf8_unicode_ci NULL DEFAULT NULL COMMENT '医院具体地址',
                `location` VARCHAR ( 100 ) CHARACTER SET utf8 COLLATE utf8_unicode_ci NULL DEFAULT NULL COMMENT '医院所在经纬度',
                `del_flag` INT NULL DEFAULT 0 COMMENT '删除状态 0正常 1已删除',
                `create_by` VARCHAR ( 32 ) CHARACTER SET utf8 COLLATE utf8_general_ci NULL DEFAULT NULL COMMENT '创建人',
                `create_time` datetime NULL DEFAULT NULL COMMENT '创建时间',
                `update_by` VARCHAR ( 32 ) CHARACTER SET utf8 COLLATE utf8_general_ci NULL DEFAULT NULL COMMENT '更新人',
                `update_time` datetime NULL DEFAULT NULL COMMENT '更新时间',
                `image` VARCHAR ( 255 ) CHARACTER SET utf8 COLLATE utf8_unicode_ci NULL DEFAULT NULL COMMENT '图片',
                PRIMARY KEY ( `id` ) USING BTREE )
                ENGINE = INNODB CHARACTER SET = utf8 COLLATE = utf8_unicode_ci COMMENT = '医院表' ROW_FORMAT = Dynamic;
                """;
        sqlEngine.builder(vectorStore,ragEngine,chatModel,resourceLoader).train(TrainBuilder.builder().content(ddl).tableName("dtp_hospital").policy(TrainPolicyType.DDL).build());
    }


    @RequestMapping("question")
    public void question(){
        String sql="SELECT * FROM dtp_hospital WHERE DISTRICT LIKE '%黄浦区%';";
        String question="在黄浦区的医院有哪些？";
        sqlEngine.builder(vectorStore,ragEngine,chatModel,resourceLoader).train(TrainBuilder.builder().content(sql).question(question).tableName("dtp_hospital").policy(TrainPolicyType.SQL).build());
    }



    @PostMapping("/chat")
    public ResponseBodyEmitter chat(@RequestBody ChatReq chatReq, HttpServletResponse servletResponse) {
        ResponseBodyEmitter emitter = new ResponseBodyEmitter();
        servletResponse.setContentType(MediaType.TEXT_EVENT_STREAM_VALUE);
        String sql = sqlEngine.setChatModel(chatModel).generateSql(chatReq);
        sql="SELECT * FROM dtp_hospital WHERE DISTRICT LIKE '%黄浦区%';";
        Object object = sqlEngine.executeSql(sql);
        JSONObject echarsJson = sqlEngine.generateEcharsJson(object.toString());
        Flux<ChatResponse> stream = sqlEngine.generateSummary(chatReq.getQuestion(), object.toString());
        stream.publishOn(Schedulers.boundedElastic())
                .doOnError(emitter::completeWithError)
                .doOnComplete(()->{
                            try {
                                if(echarsJson!=null){
                                    emitter.send("<echar>"+echarsJson+"</echar>");
                                }
                            } catch (IOException e) {
                                throw new RuntimeException(e);
                            }finally {
                                emitter.complete();
                            }
                        }
                )
                .subscribe(data -> {
                    String text = data.getResult().getOutput().getText();
                    try {
                        if(StrUtil.isNotEmpty(text)){
                            emitter.send(text);
                            //Thread.sleep(50);
                        }
                    } catch (Exception e) {
                        emitter.completeWithError(e);
                    }
                });
        return emitter;
    }


    @PostMapping("/chatV2")
    public Map chatV2(@RequestBody ChatReq chatReq) {
        sqlEngine = sqlEngine.builder(vectorStore, ragEngine, chatModel, resourceLoader);
        String sql = sqlEngine.setChatModel(chatModel).generateSql(chatReq);
        Object object = sqlEngine.executeSql(sql);
        //JSONObject echarsJson = sqlEngine.generateEcharsJson(object.toString());
        String summaryV2 = sqlEngine.generateSummaryV2(null, chatReq.getQuestion(), "", object.toString());
        Map<@Nullable String, @Nullable Object> map = Maps.newHashMap();
        //map.put("echar",echarsJson);
        map.put("sql",sql);
        map.put("chatResponse",summaryV2);
        return map;
    }


    @PostMapping("train")
    public void train(
            @RequestParam(value = "question",required = false) String  question,
            @RequestParam(value = "content",required = false) String  content,
            @RequestParam(value = "policy",required = false) String  policy,
            @RequestParam(value = "tableName",required = false) String  tableName
    ) {
        TrainBuilder builder = TrainBuilder.builder()
                .question(question)
                .content(content)
                .policy(TrainPolicyType.valueOf(policy))
                .tableName(tableName).build();
        sqlEngine = sqlEngine.builder(vectorStore, ragEngine, chatModel, resourceLoader);
        sqlEngine.train(builder);
    }



}
